import pyspiel
import numpy as np


class Env(object):
    def __init__(self, n, k):
        game = "oshi_zumo"
        game_settings = dict(coins=n, size=k)
        self.n = n
        self.k = k

        self.game = pyspiel.load_game(game, game_settings)
        self.state = None

    def init(self):
        self.state = self.game.new_initial_state()

        return self.get_obs(), self.is_done()

    def get_obs(self):
        obs = self.state.observation_tensor(0)  # 玩家1和玩家2的状态是一样的
        o = np.where(np.array(obs) == 1)[0]
        o = o - np.array([0, self.n + 1, (self.n + 1) * 2])

        return tuple(o)

    def is_done(self):
        return self.state.is_terminal()

    def get_reward(self):
        return self.state.rewards()

    def step(self, actions):
        self.state.apply_actions(actions)

        return self.get_obs(), self.is_done()
